package week07

// 208. 实现 Trie (前缀树)
// https://leetcode-cn.com/problems/implement-trie-prefix-tree/

type Node struct {
	count int
	child map[int32]*Node
}

func newNode() *Node {
	return &Node{
		count: 0,
		child: make(map[int32]*Node),
	}
}

type Trie struct {
	root *Node
}


/** Initialize your data structure here. */
func Constructor() Trie {
	return Trie{
		root: newNode(),
	}
}


/** Inserts a word into the trie. */
func (this *Trie) Insert(word string)  {
	this.find(word, true, false)
}


/** Returns if the word is in the trie. */
func (this *Trie) Search(word string) bool {
	return this.find(word, false, false)
}


/** Returns if there is any word in the trie that starts with the given prefix. */
func (this *Trie) StartsWith(prefix string) bool {
	return this.find(prefix, false, true)
}

func (this *Trie) find(word string, isInsert, isFindPrefix bool) bool {
	var curr = this.root
	for _, s := range word {
		next, ok := curr.child[s]
		if !ok {
			if !isInsert {
				return false
			}
			next = newNode()
			curr.child[s] = next
		}
		curr = next
	}
	if isFindPrefix {
		return true
	}
	if isInsert {
		curr.count++
	}
	return curr.count > 0
}


/**
 * Your Trie object will be instantiated and called as such:
 * obj := Constructor();
 * obj.Insert(word);
 * param_2 := obj.Search(word);
 * param_3 := obj.StartsWith(prefix);
 */